import requests
import json
from typing import Optional, Dict, Any, Generator


class OllamaClient:
    def __init__(self, base_url: str = "http://localhost:11434"):
        self.base_url = base_url.rstrip("/")

    def is_available(self) -> bool:
        """Check if Ollama service is available."""
        try:
            response = requests.get(f"{self.base_url}/api/tags", timeout=5)
            return response.status_code == 200
        except requests.exceptions.RequestException:
            return False

    def get_models(self) -> list:
        """Get list of available models."""
        try:
            response = requests.get(f"{self.base_url}/api/tags")
            response.raise_for_status()
            data = response.json()
            return data.get("models", [])
        except requests.exceptions.RequestException as e:
            print(f"Error fetching models: {e}")
            return []

    def translate(
        self,
        text: str,
        model: str = "llama2",
        source_lang: str = "English",
        target_lang: str = "French",
        custom_prompt: str = "",
    ) -> Optional[str]:
        """Translate text using Ollama with optional custom prompt."""

        # Build the prompt based on whether custom prompt is provided
        if custom_prompt:
            prompt = f"""Translate the following text from {source_lang} to {target_lang}.

            {custom_prompt}

Source text: {text}

You only need to return the translation only:"""
        else:
            prompt = f"""Translate the following text from {source_lang} to {target_lang}.

Source text: {text}

You only need to return the translation only:"""

        payload = {"model": model, "prompt": prompt, "stream": False}

        try:
            response = requests.post(
                f"{self.base_url}/api/generate", json=payload, timeout=30
            )
            response.raise_for_status()

            data = response.json()
            return data.get("response", "").strip()

        except requests.exceptions.RequestException as e:
            print(f"Translation error: {e}")
            return None

    def translate_stream(
        self,
        text: str,
        model: str = "llama2",
        source_lang: str = "English",
        target_lang: str = "French",
        custom_prompt: str = "",
    ) -> Generator[str, None, None]:
        """Stream translation text using Ollama."""
        
        # Build the prompt based on whether custom prompt is provided
        if custom_prompt:
            prompt = f"""Translate the following text from {source_lang} to {target_lang}.

{custom_prompt}

Source text: {text}

You only need to give the translation only:"""
        else:
            prompt = f"""Translate the following text from {source_lang} to {target_lang}.

Source text: {text}

You only need to give the translation only:"""

        payload = {
            "model": model,
            "prompt": prompt,
            "stream": True  # Enable streaming
        }

        try:
            response = requests.post(
                f"{self.base_url}/api/generate",
                json=payload,
                stream=True,
                timeout=30
            )
            response.raise_for_status()

            # Process streaming response
            for line in response.iter_lines():
                if line:
                    try:
                        data = json.loads(line.decode('utf-8'))
                        if 'response' in data:
                            yield data['response']
                        if data.get('done', False):
                            break
                    except json.JSONDecodeError:
                        continue

        except requests.exceptions.RequestException as e:
            print(f"Stream translation error: {e}")
            yield "Translation error occurred"

    def test_connection(self) -> tuple[bool, str]:
        """Test connection to Ollama service."""
        try:
            response = requests.get(f"{self.base_url}/api/tags", timeout=5)
            if response.status_code == 200:
                return True, "Connected successfully"
            else:
                return False, f"HTTP {response.status_code}: {response.text}"
        except requests.exceptions.RequestException as e:
            return False, f"Connection failed: {e}"
